import sys
from PySide6.QtWidgets import (QApplication, QWidget, QTableWidget, 
                             QTableWidgetItem, QVBoxLayout, QLabel,QHBoxLayout)
from PySide6.QtCore import Qt,QMimeData, QByteArray, QDataStream,QPoint
from PySide6.QtGui import QDrag, QColor, QPixmap, QPainter

class basetable(QTableWidget):
    def __init__(self, rows=5, cols=3, parent=None):
        super().__init__(rows, cols, parent)
        self.setSelectionMode(QTableWidget.ExtendedSelection)
        self.setDragDropMode(QTableWidget.DragDrop)
        self.setAcceptDrops(True)
        self.setDragEnabled(True)
        self.setDefaultDropAction(Qt.CopyAction)

        # 状态变量
        self.last_highlighted = None  # 记录上一个高亮的单元格
        self.forbidden_cells =set()# {(0, 0),(0, 1)}  # 禁止拖放的单元格(行,列)
        self.number_columns = set()  # 只接受数字的列(索引)
        self.cleartable()
        self.refreshBackground()
        for i in range(self.columnCount()):
            self.setColumnWidth(i,60) #设置列宽

    def refreshBackground(self):
        for row in range(self.rowCount()):
            for col in range(self.columnCount()):
                item=self.item(row,col)
                if (row, col) in self.forbidden_cells:
                    item.setBackground(QColor(255, 220, 220))
                    item.setToolTip("禁止拖放")
                else:
                    item.setBackground(QColor(255, 255, 255))
    def cleartable(self):
        for row in range(self.rowCount()):
            for col in range(self.columnCount()):
                item = QTableWidgetItem(f"")
                item.setTextAlignment(Qt.AlignCenter)
                self.setItem(row, col, item)
        return True
    
    def startDrag(self, supported_actions):
        """重写开始拖动方法，自定义拖动数据和视觉反馈"""
        selected_items = self.selectedItems()
        if not selected_items:
            return super().startDrag(supported_actions)
            # 执行拖动并处理结果

        drag=self.creatDrag(selected_items[0])
        if drag is False :
            return False
        drag.exec_(supported_actions)

    def creatDrag(self, item):
        selected_items = self.selectedItems()
        # 过滤掉禁止拖动的单元格
        if (item.row(), item.column())  in self.forbidden_cells :
            return False
        
        # 创建MIME数据
        mime_data = QMimeData()
        
        # 1. 添加Qt默认表格格式(便于表格间拖放)
        # 1. 生成application/x-qabstractitemmodeldatalist格式数据
        byte_array = QByteArray()
        data_stream = QDataStream(byte_array, QDataStream.WriteOnly)

        item=selected_items[0]
        data_stream.writeInt32(item.row())
        data_stream.writeInt32(item.column())
        
        # 写入角色-值对数量（这里包含3种角色）
        data_stream.writeInt32(4)
        
        # 角色1：显示文本
        data_stream.writeInt32(Qt.DisplayRole)
        data_stream.writeQVariant(item.text())
        
        # 角色2：背景色
        data_stream.writeInt32(Qt.BackgroundRole)
        data_stream.writeQVariant(item.background())
        
        # 角色3：提示文本
        data_stream.writeInt32(Qt.ToolTipRole)
        data_stream.writeQVariant(item.toolTip())
        #角色4：对齐方式
        data_stream.writeInt32(Qt.TextAlignmentRole)
        data_stream.writeQVariant(item.textAlignment())
        
        mime_data.setData("application/x-qabstractitemmodeldatalist", byte_array)
        
        # 2. 添加文本格式(兼容其他控件)
        text_data = f"{item.row()},{item.column()},{item.text()}"
        mime_data.setText(text_data)
        
        # 创建拖动对象
        drag = QDrag(self)
        drag.setMimeData(mime_data)
        
        # 设置拖动图标
        drag.setPixmap(self._create_drag_pixmap(item.text()))
        drag.setHotSpot(QPoint(-10, 0))

        return drag
        


    def _create_drag_pixmap(self, dataText):
        """创建自定义拖动图标"""
        pixmap = QPixmap(80, 25)
        pixmap.fill(QColor(255, 255, 255,200))
        
        painter = QPainter(pixmap)
        painter.setPen(QColor(0, 0, 0))
        painter.drawText(pixmap.rect(), Qt.AlignCenter, f"{dataText}")
        painter.drawRect(0, 0, pixmap.width()-1, pixmap.height()-1)
        painter.end()
        
        return pixmap
    
    def dragEnterEvent(self, event):
        """拖入时验证数据"""
        if (
            event.mimeData().hasText() or 
            event.mimeData().hasFormat("application/x-qabstractitemmodeldatalist")
        ):
            event.acceptProposedAction()
        else:
            event.ignore()

    def dragMoveEvent(self, event):
        """拖动过程中提供视觉反馈和位置验证"""
        # 清除上一次高亮
        self._clear_highlight()
        
        # 获取目标位置
        pos = event.pos()
        index = self.indexAt(pos)
        
        if not index.isValid():
            event.ignore()
            return
            
        row, col = index.row(), index.column()
        
        # 检查是否为禁止区域
        if (row, col) in self.forbidden_cells:
            event.ignore()
            return
            
        # 高亮目标单元格
        self._highlight_cell(row, col)
        
        event.acceptProposedAction()

    def dragLeaveEvent(self, event):
        """拖动离开时清除高亮"""
        self._clear_highlight()
        super().dragLeaveEvent(event)

    def dropEvent(self, event):
        """处理最终放下的数据"""
        # 清除高亮
        self._clear_highlight()
        
        # 获取目标位置
        pos = event.pos()
        index = self.indexAt(pos)
        target_row, target_col = 0, 0
        
        if index.isValid():
            target_row, target_col = index.row(), index.column()
        else:
            # 拖到空白处，放到最后一行
            target_row = self.rowCount()
            target_col = 0
            self.insertRow(target_row)
        
        # 解析并处理数据
        mime_data = event.mimeData()
        if mime_data.hasFormat("application/x-qabstractitemmodeldatalist"):
            self._process_model_data(mime_data, target_row, target_col)
        elif mime_data.hasText():
            self._process_text_data(mime_data.text(), target_row, target_col)
            
        event.acceptProposedAction()
# ------------------------------
    # 辅助方法
    # ------------------------------
    def _highlight_cell(self, row, col):
        """高亮目标单元格"""
        item = self.item(row, col) or QTableWidgetItem()
        self.last_highlighted = {
            "row": row,
            "col": col,
            "bg": item.background()
        }
        item.setBackground(QColor(255, 245, 157))
        if not self.item(row, col):
            self.setItem(row, col, item)

    def _clear_highlight(self):
        """清除单元格高亮"""
        if self.last_highlighted:
            row = self.last_highlighted["row"]
            col = self.last_highlighted["col"]
            item = self.item(row, col)
            if item:
                item.setBackground(self.last_highlighted["bg"])
            self.last_highlighted = None
    def _process_text_data(self, text, start_row, start_col):
        """处理文本格式数据"""
        items = text.split(';')
        current_row, current_col = start_row, start_col
        
        for item_data in items:
            parts = item_data.split(',', 2)
            if len(parts) != 3:
                continue
            _, _, text = parts
            
            # 检查数字列
            if current_col in self.number_columns and not text.replace('.', '', 1).isdigit():
                text = f"[{text}]"  # 标记无效数字
            
            self._insert_item(current_row, current_col, text)
            
            # 处理列溢出
            current_col += 1
            if current_col >= self.columnCount():
                current_row += 1
                current_col = 0
                if current_row >= self.rowCount():
                    self.insertRow(current_row)

    def _process_model_data(self, mime_data, start_row, start_col):
        """处理Qt表格格式数据"""
        data = mime_data.data("application/x-qabstractitemmodeldatalist")
        stream = QDataStream(data, QDataStream.ReadOnly)
        
        row_s = stream.readInt32()
        current_row, current_col = start_row, start_col
        

        # 源行列索引(这里不需要使用)
        col_s=stream.readInt32()
        item_count=stream.readInt32()
        item_data = {}
        
        for _ in range(item_count):
            role = stream.readInt32()
            value = stream.readQVariant()
            item_data[role] = value
        
        # 提取显示文本和样式
        text = item_data.get(Qt.DisplayRole, "未知数据")
        background = item_data.get(Qt.BackgroundRole, QColor(255, 255, 255))
        alignCenter =item_data.get(Qt.TextAlignmentRole,Qt.AlignCenter)
        item = QTableWidgetItem(text)
        item.setTextAlignment(alignCenter)  
        # 插入单元格
        self.setItem(current_row, current_col, item)
        # item.setBackground(background)

class TeacherList(basetable):
    def __init__(self, rows=100, cols=5, parent=None):
        super().__init__(rows, cols, parent)
        self.setHorizontalHeaderLabels(["姓名","倒1","倒2","总计","排序"])
        self.setColumnWidth(0,60)
        self.setColumnWidth(1,30)
        self.setColumnWidth(2,30)
        self.setColumnWidth(3,40)
        self.setSortingEnabled(True)
        self.setDragDropMode(QTableWidget.DragOnly)

class TeacherExchange(basetable):
    def __init__(self,examRoomNum=20,examTeacherNum=2,examNum=5, parent=None):
        
        self.examRoomNum=examRoomNum
        self.examTeacherNum=examTeacherNum
        self.examNum=examNum

        super().__init__(examRoomNum+1, examTeacherNum*examNum+1, parent)
        self.forbidden_first_row_col={(0,i*examTeacherNum+1) for i in range(examNum)}
        self.forbidden_first_row_col.update({(i,0) for i in range(examRoomNum+1)})

        for i in range(examNum):
            self.setSpan(0,i*examTeacherNum+1,1,examTeacherNum)

        self.forbidden_cells.update(self.forbidden_first_row_col)
        self.setRowHeight(0,100)
        self.horizontalHeader().setVisible(False)
        self.verticalHeader().setVisible(False)
        self.refreshBackground()

    def refreshBackground(self):
        for row in range(self.rowCount()):
            for col in range(self.columnCount()):
                item=self.item(row,col)
                if (row, col) in self.forbidden_cells:
                    if(row,col) in self.forbidden_first_row_col:
                        item.setBackground(QColor("#FAEBD7"))
                    else:
                        item.setBackground(QColor(255, 220, 220))
                    item.setToolTip("禁止拖放")
                else:
                    item.setBackground(QColor(255, 255, 255))


class DemoWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("QTableWidget 重写拖放方法示例")
        self.setGeometry(100, 100, 1000, 600)
        
        layout = QVBoxLayout()
        
        info_label = QLabel("""
        拖放功能说明：<br>
        - 红色单元格：禁止拖放<br>
        - 绿色单元格：仅接受数字（非数字会标记）<br>
        - 拖动时会高亮目标位置<br>
        - 支持表格间互拖和表格内拖放
        """)
        layout.addWidget(info_label)
        
        # 创建两个表格用于测试
        table_layout = QHBoxLayout()
        self.table1 = basetable()
        self.table2 = TeacherExchange()
        self.table3 = TeacherList()
        self.table3.setDragDropMode(QTableWidget.DragOnly)
        table_layout.addWidget(self.table1)
        table_layout.addWidget(self.table2)
        table_layout.addWidget(self.table3)
        
        layout.addLayout(table_layout)
        self.setLayout(layout)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = DemoWindow()
    window.show()
    sys.exit(app.exec())
    